import numpy as np
from keras import Sequential
from keras.layers import Dense
import matplotlib.pyplot as plt

x_train = np.linspace(-2, 6, 200)
np.random.shuffle(x_train)
y_train = x_train * 0.5 + 2 + 0.1 * np.random.randn(200)

x_test = x_train[150:, ]
y_test = y_train[150:, ]
x_train = x_train[:150, ]
y_train = y_train[:150, ]

model = Sequential()
model.add(Dense(1, input_shape=(1,)))

model.compile(optimizer="sgd", loss="mse")

epoch = 1000
split1 = 0.009

for i in range(epoch):
    cost = model.train_on_batch(x_train, y_train)
    if i % 20:
        print(cost)

w, b = model.layers[0].get_weights()
print("w:", w)
print("b:", b)

y_pre = model.predict(x_test)
plt.plot(x_test, y_pre)
plt.scatter(x_test, y_test)
plt.show()
